#include "wrapper_sstructmg.h"

typedef int idx_t;
typedef double data_t;

// const idx_t dims_per_blk[] = {
//     40 , 15 , 25 ,
//     192, 20 , 40 ,
//     25 , 65 , 40 ,
//      4 , 25 , 40 ,
//      2 ,  9 ,  8 ,
//     40 , 25 ,  8 ,
//     40 , 25 , 25 ,
//     40 , 25 , 35 ,
//     15 ,  5 , 40 ,
//      5 ,  3 , 40 ,
//      3 , 15 , 40 ,
//      2 ,  9 , 40 ,
//     40 , 25 , 40 ,
// };
// const idx_t offd_nnz [] = {
// //    0      1       2       3       4       5       6       7       8       9       10      11      12
//       0,    600,   1000,      0,      0,      0,      0,      0,      0,      0,      0,      0,      0, // 0
//     600,      0,   2600,    160,      0,    320,   1000,   1400,      0,      0,      0,      0,   1600, // 1
//    1000,   2600,      0,      0,      0,      0,   1000,      0,      0,      0,      0,      0,      0, // 2
//       0,    160,      0,      0,      0,   1000,      0,   1000,      0,      0,      0,      0,      0, // 3
//       0,      0,      0,      0,      0,     72,      0,      0,      0,      0,      0,     18,      0, // 4
//       0,    320,      0,   1000,     72,      0,      0,      0,      0,      0,      0,      0,   1000, // 5
//       0,   1000,   1000,      0,      0,      0,      0,   1000,      0,      0,      0,      0,      0, // 6
//       0,   1400,      0,   1000,      0,      0,   1000,      0,      0,      0,      0,      0,      0, // 7
//       0,      0,      0,      0,      0,      0,      0,      0,      0,    200,    600,      0,      0, // 8
//       0,      0,      0,      0,      0,      0,      0,      0,    200,      0,    120,      0,    200, // 9
//       0,      0,      0,      0,      0,      0,      0,      0,    600,    120,      0,      0,    600, // 10
//       0,      0,      0,      0,     18,      0,      0,      0,      0,      0,      0,      0,    360, // 11
//       0,   1600,      0,      0,      0,   1000,      0,      0,      0,    200,    600,    360,      0, // 12
// };

// const idx_t dims_per_blk[] = {
//     //  80,  25,   50,
//     // 374,  35,   80,
//     //  50, 130,   80,
//     //   8,  50,   80,
//     //   4,  19,   16,
//     //  80,  50,   16,
//     //  80,  50,   50,
//     //  80,  50,   70,
//     //  29,  11,   75,
//     //  11,   5,   75,
//     //   5,  29,   75,
//     //   4,  19,   75,
//     //  75,  50,   80,
//  25,  50, 80,
//  35, 374, 80,
// 130,  50, 80,
//   8,  50, 80,
//  16,   4, 19,
//  50,  16, 80,
//  50,  80, 50,
//  70,  80, 50,
//  75,  29, 11,
//  75,   5, 11,
//  75,  29, 5,
//  75,   4, 19,
//  75,  50, 80,
// };
// const idx_t offd_nnz [] = {
//       0,    2000,    4000,       0,       0,       0,       0,       0,       0,       0,       0,       0,       0, //  6000
//    2000,       0,   10400,     640,       0,    1280,    4000,    5600,       0,       0,       0,       0,    6000, // 29920
//    4000,   10400,       0,       0,       0,       0,    4000,       0,       0,       0,       0,       0,       0, // 18400
//       0,     640,       0,       0,       0,    4000,       0,    4000,       0,       0,       0,       0,       0, //  8640
//       0,       0,       0,       0,       0,     304,       0,       0,       0,       0,       0,      76,       0, //   380
//       0,    1280,       0,    4000,     304,       0,       0,       0,       0,       0,       0,       0,    4000, //  9584
//       0,    4000,    4000,       0,       0,       0,       0,    4000,       0,       0,       0,       0,       0, // 12000
//       0,    5600,       0,    4000,       0,       0,    4000,       0,       0,       0,       0,       0,       0, // 13600
//       0,       0,       0,       0,       0,       0,       0,       0,       0,     825,    2175,       0,       0, //  3000
//       0,       0,       0,       0,       0,       0,       0,       0,     825,       0,     375,       0,     825, //  2025
//       0,       0,       0,       0,       0,       0,       0,       0,    2175,     375,       0,       0,    2175, //  4725
//       0,       0,       0,       0,      76,       0,       0,       0,       0,       0,       0,       0,    1425, //  1501
//       0,    6000,       0,       0,       0,    4000,       0,       0,       0,     825,    2175,    1425,       0, // 14425
// };

const idx_t dims_per_blk[] = {
230, 60, 108,
230,60, 108,
104, 24, 5,
39,	15,	80,
15,	5,	80,
5,	39,	80,
460, 150, 108,
};
const idx_t offd_nnz [] = {
      0,    6480,       0,       0,       0,       0,   24840, // 31320
   6480,       0,    2496,       0,    1200,    3120,   24840, // 38136
      0,    2496,       0,       0,       0,       0,       0, //  2496
      0,       0,       0,       0,    1200,    3120,       0, //  4320
      0,    1200,       0,    1200,       0,     400,       0, //  2800
      0,    3120,       0,    3120,     400,       0,       0, //  6640
  24840,   24840,       0,       0,       0,       0,       0, // 49680
};

// const std::string data_path = 
// #ifdef __aarch64__
// "/storage/hpcauser/zongyi/SStructMG-beta/data/kvlcc2/280w";
// #elif defined(__x86_64__)
// "/home/export/base/ycsc_xuew/xuew/online1/zongyi/SStructMG/data/kvlcc2";
// #endif

#include <unordered_set>

int main(int argc, char* argv[])
{
    int my_pid, num_procs;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &my_pid);
    MPI_Comm_size(MPI_COMM_WORLD, &num_procs);
    if (my_pid == 0) printf("\033[1;35mMax threads %d\033[0m \n", omp_get_max_threads());

    idx_t arg_cnt = 1;
    const std::string data_path(argv[arg_cnt++]);
    const std::string assign_file(argv[arg_cnt++]);
    // const std::string pathname(argv[arg_cnt++]);
    const std::string its_name = std::string(argv[arg_cnt++]),
                    prc_name = std::string(argv[arg_cnt++]);
    std::string config_mg_file;
	if (prc_name == "GMG") config_mg_file = std::string(argv[arg_cnt++]);

    const idx_t glb_nparts = sizeof(dims_per_blk) / sizeof(idx_t) / 3;
    idx_t blk_offsets[glb_nparts + 1];
    blk_offsets[0] = 0;
    for (idx_t b = 0; b < glb_nparts; b++) {
        blk_offsets[b + 1] = blk_offsets[b] 
            + dims_per_blk[b*3  ] * dims_per_blk[b*3+1] * dims_per_blk[b*3+2];
    }
    // const idx_t tot_nrows = blk_offsets[glb_nparts];

    SStructMG_Init();
    {
        SStructGrid<idx_t> ssgrid(MPI_COMM_WORLD, 3, glb_nparts);
        const idx_t halos[3] = {1, 1, 1};
        {// 添加本进程所拥有的块和盒子
            FILE * fp = fopen(assign_file.c_str(), "r");
            idx_t blk_id, blk_exts[3], blk_nprocs;
            const bool periods[3] = {false, false, false};
            std::unordered_set<idx_t> my_added_blks;

            idx_t tmp;

                fscanf(fp, "%d", &tmp); assert(tmp == glb_nparts);
                for (idx_t b  = 0; b < glb_nparts; b++) {
                    fscanf(fp, "%d (%d,%d,%d) %d", &blk_id, &blk_exts[0], &blk_exts[1], &blk_exts[2], &blk_nprocs);
#if 0
                    printf("Proc %d read Part %d: (%d,%d,%d) %d\n", my_pid, blk_id,
                        blk_exts[0], blk_exts[1], blk_exts[2], blk_nprocs);
#endif
                    assert(b == blk_id);
                    for (idx_t p = 0; p < blk_nprocs; p++) {// 读取这么多次
                        idx_t pid, begs[3], ends[3];
                        fscanf(fp, "%d (%d,%d,%d) (%d,%d,%d)", &pid, &begs[0], &begs[1], &begs[2], &ends[0], &ends[1], &ends[2]);
#if 0
                        printf("Proc %d read Box %d: %d (%d,%d,%d) (%d,%d,%d)\n", my_pid, p, pid,
                            begs[0], begs[1], begs[2], ends[0], ends[1], ends[2] );
#endif                        
                        if (pid == my_pid) {// 是本进程负责的盒子
                            if (my_added_blks.find(blk_id) == my_added_blks.end()) {// 本进程之前没有负责过这个块
                                ssgrid.addBlock(blk_id, blk_exts, periods);
                                my_added_blks.insert(blk_id);// 记录一下已添加该块
                            }
                            ssgrid.addBox(blk_id, begs, ends);
#if 0
                            printf("Proc %d addBox to Blk %d (%d,%d,%d) ~ (%d,%d,%d)\n", my_pid, blk_id,
                                begs[0], begs[1], begs[2], ends[0], ends[1], ends[2]);
#endif
                        }// 否则pass到下一个
                    }
                }
            
            fclose(fp);
        }
        
        ssgrid.assemble();
        const idx_t my_nblks = ssgrid.my_blk_list.size();
        idx_t shalos[my_nblks * 3];
        for (idx_t k = 0; k < my_nblks; k++) {
            shalos[k * 3 + 0] = halos[0];
            shalos[k * 3 + 1] = halos[1];
            shalos[k * 3 + 2] = halos[2];
        }

        // 根据半结构网格建立向量
        par_SstructVector<idx_t, data_t> b(ssgrid, shalos, true);
        par_SstructVector<idx_t, data_t> x(b), y(b), x0(b);
        {// 读入向量数据
            for (idx_t blk = 0, ptr = 0; blk < glb_nparts; blk++) {
                if (ptr < my_nblks && ssgrid.my_blk_list[ptr].blk_id == blk) {
                    // printf("Proc %d reading blk %d\n", my_pid, blk);
                    b.struct_vectors[ptr]->IO_data(0, data_path, "b.B" + std::to_string(blk));
                    x0.struct_vectors[ptr]->IO_data(0, data_path, "x0.B"+ std::to_string(blk));
                    ptr++;
                }
                // MPI_Barrier(MPI_COMM_WORLD);
                // if (my_pid == 0) printf("Load %d done\n", blk);
            }
            y.set_val(0.0, true);
        }
        
        // 根据半结构网格建立矩阵
        __int128_t masks[my_nblks]; for (idx_t k = 0; k < my_nblks; k++) masks[k] = stencil_mask_3d7;
        par_SstructMatrix<idx_t, data_t, data_t> A(ssgrid, shalos, masks);
        int proc_diag_nrows = 0, proc_offd_nnz = 0;
        {// 读入矩阵数据
            // 首先处理对角块，也即结构化部分
            for (idx_t blk = 0, ptr = 0; blk < glb_nparts; blk++) {
                if (ptr < my_nblks && ssgrid.my_blk_list[ptr].blk_id == blk) {
                    idx_t nrows = 0;
                    for (idx_t bx = 0; bx < (idx_t) ssgrid.my_blk_list[ptr].my_box_list.size(); bx++) {
                        const Box<idx_t> & box = ssgrid.my_blk_list[ptr].my_box_list[bx];
                        nrows +=  (box.ends[0] - box.begs[0])
                                * (box.ends[1] - box.begs[1])
                                * (box.ends[2] - box.begs[2]);
                    }
                    proc_diag_nrows += nrows;
                    A.diag_matrixs[ptr]->IO_data(0, data_path, "B" + std::to_string(blk) + "_Diag");
                    ptr++;
                }
                // MPI_Barrier(MPI_COMM_WORLD);
                // if (my_pid == 0) printf("Load %d done\n", blk);
            }
            
            FILE * fp = nullptr;
            // 然后处理非对角快，也即非结构部分
            A.init_offd();
            std::vector< Entry<idx_t, data_t> > all_nnz;
            for (idx_t mk = 0; mk < my_nblks; mk++) {// 遍历每一个本进程需要参与的块
                const Block<idx_t> & blk = ssgrid.my_blk_list[mk];
                const idx_t my_nboxs = blk.my_box_list.size();
                for (idx_t nbk = 0; nbk < glb_nparts; nbk++) {// 遍历每一个可能的邻居块
                    const idx_t nnz = offd_nnz[blk.blk_id * glb_nparts + nbk];
                    if (nnz == 0) continue;
                    // 如果该邻居块对本块有影响
                    const idx_t ngb_dims[3] = { dims_per_blk[nbk * 3 + 0],
                                                dims_per_blk[nbk * 3 + 1],
                                                dims_per_blk[nbk * 3 + 2] };
                    // 读入该非对角块的数据
                    COO<idx_t, data_t> * buf = new COO<idx_t, data_t> [nnz];
                    const std::string filename = data_path + "/B" + std::to_string(blk.blk_id)+"_Offd"+std::to_string(nbk);
                    fp = fopen(filename.c_str(), "rb");
                    if (fp == nullptr) {
                        printf("Proc %d open %s failed\n", my_pid, filename.c_str());
                        MPI_Abort(MPI_COMM_WORLD, -2);
                    }
                    size_t ret = fread(buf, sizeof(COO<idx_t, data_t>), nnz, fp); assert(ret == (size_t) nnz);
                    fclose(fp);

                    StructDescriptor<idx_t> row, col;
                    row.id = blk.blk_id;
                    col.id = nbk;
                    for (idx_t k = 0; k < nnz; k++) {// 遍历每一个可能的非零元
                        const idx_t loc_1D = buf[k].r - blk_offsets[blk.blk_id],// 该非零元的行序号在本块内的一维索引
                                    ngb_1D = buf[k].c - blk_offsets[nbk];// 该非零元的列序号在邻居块内的一维索引
                        idx_t * index = row.index;// 该非零元的行在本块内的三维坐标
                        index[0] =  loc_1D / (blk.extents[1] * blk.extents[2]);
                        index[1] = (loc_1D - index[0] * blk.extents[1] * blk.extents[2]) / blk.extents[2];
                        index[2] =  loc_1D - index[0] * blk.extents[1] * blk.extents[2] - index[1] * blk.extents[2];
                        idx_t * to_id = col.index;// 该非零元的列在邻居块内的三维坐标
                        to_id[0] =  ngb_1D / (ngb_dims[1] * ngb_dims[2]);
                        to_id[1] = (ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2]) / ngb_dims[2];
                        to_id[2] =  ngb_1D - to_id[0] * ngb_dims[1] * ngb_dims[2] - to_id[1] * ngb_dims[2];

                        for (idx_t bx = 0; bx < my_nboxs; bx++) {
                            const idx_t * begs = blk.my_box_list[bx].begs,
                                        * ends = blk.my_box_list[bx].ends;
                            if (begs[0] <= index[0] && index[0] < ends[0] &&
                                begs[1] <= index[1] && index[1] < ends[1] &&
                                begs[2] <= index[2] && index[2] < ends[2] )
                            {// 该行属于这个盒子
                                // if (row.id == 4) {
                                //     printf("Proc %d add %d (%d,%d,%d) -- %d (%d,%d,%d) %.3e\n", my_pid,
                                //         row.id, row.index[0], row.index[1], row.index[2],
                                //         col.id, col.index[0], col.index[1], col.index[2], buf[k].v);
                                // }
                                all_nnz.push_back(Entry<idx_t, data_t>(row, col, buf[k].v));
                                break;
                            }
                        }// 找到所属盒子后跳出循环
                    }
                    delete [] buf;
                }
            }
            proc_offd_nnz = all_nnz.size();
            // 是否需要混洗
            unsigned seed = std::chrono::system_clock::now().time_since_epoch().count();
            std::shuffle(all_nnz.begin(), all_nnz.end(), std::default_random_engine(seed));

            A.add_entries(all_nnz.size(), all_nnz.data());
        }
        A.assemble();

        const int test_cnt = 10;
        std::vector<TEST_RECORD> records;
        for (int test = 0; test < test_cnt; test++) {
            vec_copy(x0, x);

            double ck = vec_dot<idx_t, data_t, double>(b, b);
            if (my_pid == 0) printf("(  b,   b) = %.15e\n", ck);

            ck = vec_dot<idx_t, data_t, double>(x, x);
            if (my_pid == 0) printf("(  x,   x) = %.15e\n", ck);

            A.Mult(b, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*b, A*b) = %.15e\n", ck);

            A.Mult(x, y, false);
            ck = vec_dot<idx_t, data_t, double>(y, y);
            if (my_pid == 0) printf("(A*x, A*x) = %.15e\n", ck);

            int glb_sum [2], loc_sum[2] = {proc_diag_nrows, proc_offd_nnz};
            MPI_Reduce(loc_sum, glb_sum, 2, MPI_INT, MPI_SUM, 0, MPI_COMM_WORLD);
            if (my_pid == 0) printf("Glb %d %d\n", glb_sum[0], glb_sum[1]);

            // int tmp [2] = {proc_diag_nrows, proc_offd_nnz};
            // int all_offd_nnz [num_procs * 2];
            // MPI_Allgather(tmp, 2, MPI_INT, all_offd_nnz, 2, MPI_INT, MPI_COMM_WORLD);
            // if (my_pid == 0) {
            //     int max_nrows = 0, max_offd_nnz = 0;
            //     for (int p = 0; p < num_procs; p++) {
            //         max_nrows = std::max(max_nrows, all_offd_nnz[p*2]);
            //         max_offd_nnz = std::max(max_offd_nnz, all_offd_nnz[p*2+1]);
            //         printf("Proc %d Diag nrows %d Offd nnz %d\n", p, all_offd_nnz[p*2], all_offd_nnz[p*2+1]);
            //     }
            //     printf("Max nrows %d offdnnz %d\n", max_nrows, max_offd_nnz);
            //     fflush(stdout);
            // }
            // MPI_Barrier(MPI_COMM_WORLD);

            TEST_CONFIG config;
            config.my_nblks = my_nblks;
            config.config_mg_file = config_mg_file;
            config.ssgrid_ptr = & ssgrid;
            config.rtol = 1.0e-3;
            TEST_RECORD rec;

            buildup_solver(its_name, prc_name, config);
            setup_and_solve(A, b, x, rec);

            data_t true_r_norm, b_norm;
            check_residual(A, x, b, y, true_r_norm, b_norm);

            if (my_pid == 0) {
                printf("\033[1;35mtrue ||r|| = %20.16e ||r||/||b||= %20.16e\033[0m\n", 
                    true_r_norm, true_r_norm / b_norm);
                printf("Proc %d Setup, Solve costs %.6f %.6f s\n",
                    my_pid, rec.setup, rec.solve);
            }

            stat_part_times(rec);
            records.push_back(rec);

            destroy_solver();

            fflush(stdout);
            MPI_Barrier(MPI_COMM_WORLD);
        }// test loop

        if (my_pid == 0) {// 输出最佳时间
            stat_multi_runs(records);
        }
    }// all Semi-** object must be destroyed before MPI_Finalize()
        
    SStructMG_Finalize();
    MPI_Finalize();
    return 0;
}
